{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e503765d-49ed-42f5-9bff-9bed92c6859e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#从数据集中分离出有对应标签和没有对应标签的图片数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "17b83608-39d1-4bc3-8355-4f5e8c4d5212",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "from tqdm import tqdm\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "b7d13c6a-286d-471b-9016-f68922801c1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 数据集路径\n",
    "dataset_dir = r'D:\\workspace\\datasets\\Review\\0424\\0424'  # 你的YOLO数据集路径\n",
    "images_dir = os.path.join(dataset_dir, 'images')\n",
    "labels_dir = os.path.join(dataset_dir, 'labels')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "b974cf54-4a23-402a-b55d-4cfa6dca3a45",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 0 --> 溢胶\n",
    "# 1 --> 线头\n",
    "# 2 --> 脏污\n",
    "target_class_ids = [1] "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "4e78bac2-acc7-4469-beae-ae523af1ecd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 保存筛选图片的目录\n",
    "filtered_save_dir = r'D:\\workspace\\datasets\\Review\\0424\\filtered_images_id' + \"_\".join(map(str, target_class_ids))\n",
    "no_label_save_dir = r'D:\\workspace\\datasets\\Review\\0424\\no_label_images_id' + \"_\".join(map(str, target_class_ids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "54834cd5-9d01-4e4a-a4f1-ad33a72a7cc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.makedirs(filtered_save_dir, exist_ok=True)\n",
    "os.makedirs(os.path.join(filtered_save_dir, 'labels'), exist_ok=True)\n",
    "os.makedirs(os.path.join(filtered_save_dir, 'images'), exist_ok=True)\n",
    "\n",
    "os.makedirs(no_label_save_dir, exist_ok=True)\n",
    "os.makedirs(os.path.join(no_label_save_dir, 'images'), exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "84ed1a28-ab99-405c-ae46-f676c2c5dc9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_suffixes = ('.jpg', '.jpeg', '.png', '.bmp')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "c751672f-28f5-448b-8e90-9fc71aa4d97d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🔍 总共找到 1000 张图片，开始筛选...\n"
     ]
    }
   ],
   "source": [
    "# 获取所有图片名\n",
    "image_files = [f for f in os.listdir(images_dir) if f.lower().endswith(img_suffixes)]\n",
    "\n",
    "print(f\"🔍 总共找到 {len(image_files)} 张图片，开始筛选...\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "6c3a510b-7d26-4bbe-9be0-0309369880c1",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:43<00:00, 23.13it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ 筛选完成！\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "for img_file in tqdm(image_files):\n",
    "    img_name, _ = os.path.splitext(img_file)\n",
    "    label_path = os.path.join(labels_dir, img_name + '.txt')\n",
    "    img_path = os.path.join(images_dir, img_file)\n",
    "\n",
    "    if os.path.exists(label_path):\n",
    "        with open(label_path, 'r') as f:\n",
    "            lines = f.readlines()\n",
    "\n",
    "        # 检查是否包含目标类别\n",
    "        has_target_class = False\n",
    "        for line in lines:\n",
    "            class_id = int(line.strip().split()[0])\n",
    "            if class_id in target_class_ids:\n",
    "                has_target_class = True\n",
    "                break\n",
    "\n",
    "        if has_target_class:\n",
    "            # 复制图片和标签到 filtered 目录\n",
    "            shutil.copy(img_path, os.path.join(filtered_save_dir, 'images', img_file))\n",
    "            shutil.copy(label_path, os.path.join(filtered_save_dir, 'labels', img_name + '.txt'))\n",
    "        else:\n",
    "            # 没有目标类别，复制图片到 no_label 目录\n",
    "            shutil.copy(img_path, os.path.join(no_label_save_dir, 'images', img_file))\n",
    "    else:\n",
    "        # 没有标签文件，复制图片到 no_label 目录\n",
    "        shutil.copy(img_path, os.path.join(no_label_save_dir, 'images', img_file))\n",
    "\n",
    "print(\"✅ 筛选完成！\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a55a1be-cb3c-4ff6-a6c4-9da88a91f8ed",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
